import os
from langchain.chat_models import init_chat_model
from langchain_core.documents import Document
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

#  从环境变量中读取DeepSeek的API Key
key = os.getenv("DS_API_KEY")
# print(key)
api_key = str(key)

model = init_chat_model(
    model="deepseek-chat",
    base_url="https://api.deepseek.com",
    api_key=api_key
)

# 文档对象
documents = [
    Document(
        page_content="Dogs are great companions, known for their loyalty and friendliness.",
        metadata={"source": "mammal-pets-doc"},
    ),
    Document(
        page_content="Cats are independent pets that often enjoy their own space.",
        metadata={"source": "mammal-pets-doc"},
    ),
]

file_path = r"E:\2025_File\图神经网络文献\Dynamic adaptive spatio-temporal graph neural network\Dynamic adaptive spatio-temporal graph neural network.pdf"
loader = PyPDFLoader(file_path)

docs = loader.load()

print(len(docs))
print(f"{docs[0].page_content[:200]}\n")
print(docs[0].metadata)

# 切分文档
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200, add_start_index=True
)
#  chunk_overlap 是切分的重叠  防止把语义切断\
all_splits = text_splitter.split_documents(docs)

print(len(all_splits))

# 使用huggingface上开源的embedding模型也进行文本向量化
vector_1 = embeddings.embed_query(all_splits[0].page_content)
vector_2 = embeddings.embed_query(all_splits[1].page_content)

assert len(vector_1) == len(vector_2)
print(f"Generated vectors of length {len(vector_1)}\n")
print(vector_1[:10])

# 接下来 需要将这些embedding之后的向量 进行存储  目标是为了要进行相似度检索
vector_store = InMemoryVectorStore(embeddings)
# 将文档添加到向量存储中 并返回文档的ID
ids = vector_store.add_documents(documents=all_splits)

# print(ids)
results = vector_store.similarity_search("DSTGNN")
print(results[0])  # 返回查询到的结果

# Note that providers implement different scores; the score here
# is a distance metric that varies inversely with similarity.
# 这是返回了一个分数, 衡量了相似度
results = vector_store.similarity_search_with_score("What was Nike's revenue in 2023?")
doc, score = results[0]
print(f"Score: {score}\n")
print(doc)

# Retrievers

from typing import List
from langchain_core.documents import Document
from langchain_core.runnables import chain


@chain
def retriever(query: str) -> List[Document]:
    return vector_store.similarity_search(query, k=1)


print(retriever.batch(
    [
        "DASTGNN",
        "Graph WaveNet"
    ],
))
retriever = vector_store.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 1},
)
print(retriever.batch(
    [
        "DASTGNN",
        "Graph WaveNet"
    ],
))
